import torch
from torch import nn

import torch
import torch.nn as nn

import torch
import torch.nn as nn


class SelfAttention(nn.Module):
    def __init__(self, in_channels):
        super(SelfAttention, self).__init__()
        self.query = nn.Conv2d(in_channels, in_channels , 1)
        self.key = nn.Conv2d(in_channels, in_channels , 1)
        self.value = nn.Conv2d(in_channels, in_channels, 1)
        self.gamma = nn.Parameter(torch.zeros(1))
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        batch_size, channels, height, width = x.size()

        # Compute query, key, and value embeddings
        query = self.query(x).view(batch_size, -1, height * width).permute(0, 2, 1)
        key = self.key(x).view(batch_size, -1, height * width)
        value = self.value(x).view(batch_size, -1, height * width)

        # Compute attention weights
        attention = torch.bmm(query, key)
        attention = self.softmax(attention)

        # Compute attended values and reshape
        attended = torch.bmm(value, attention.permute(0, 2, 1))
        attended = attended.view(batch_size, channels, height, width)

        # Apply gamma and add to input
        output = self.gamma * attended + x

        return output


class ASPPwithSelfAttention(nn.Module):
    def __init__(self, in_channels, out_channels, rates=[6, 12, 18]):
        super(ASPPwithSelfAttention, self).__init__()

        # ASPP convolutions with different rates
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=3, dilation=rates[0], padding=rates[0])
        self.conv3 = nn.Conv2d(in_channels, out_channels, kernel_size=3, dilation=rates[1], padding=rates[1])
        self.conv4 = nn.Conv2d(in_channels, out_channels, kernel_size=3, dilation=rates[2], padding=rates[2])

        # Self-attention
        self.attention = SelfAttention(out_channels * 4)

        # Final convolution to merge ASPP features
        self.conv_final = nn.Conv2d(out_channels * 4, out_channels, kernel_size=1)

    def forward(self, x):
        conv1 = self.conv1(x)
        print("conv1 ",conv1.shape)
        conv2 = self.conv2(x)
        conv3 = self.conv3(x)
        conv4 = self.conv4(x)

        # ASPP features
        aspp = torch.cat([conv1, conv2, conv3, conv4], dim=1)
        print("aspp",aspp.shape)
        # Apply self-attention to ASPP features
        aspp_attended = self.attention(aspp)
        print("aspp_attended",aspp_attended.shape)
        # Concatenate ASPP and self-attended ASPP features
        concatenated = torch.cat([aspp_attended, x], dim=1)

        # Apply final convolution to merge features
     #   output = self.conv_final(concatenated)

        return concatenated


img = torch.randn(2, 15, 32, 32)

model = nn.Sequential(
    ASPPwithSelfAttention(15,15)
)

model(img) # (2, 1000)